import torch
import torch.nn as nn
import nvtx

class NVTXModuleWrapper():
    def __init__(self, module: nn.Module, message=None,
                 forward_color='green', backward_color='red'):
        self.module = module
        self.message = message if message else module.__class__.__name__
        self.forward_color = forward_color
        self.backward_color = backward_color

        self.forward_rng = None
        self.backward_rng = None

    def register(self):
        self.module.register_forward_pre_hook(self._forward_pre_hook_fn)
        self.module.register_forward_hook(self._forward_post_hook_fn)
        # self.module.register_full_backward_pre_hook(self._backward_pre_hook_fn)
        # self.module.register_full_backward_hook(self._backward_post_hook_fn)

    def _forward_pre_hook_fn(self, module, input):
        self.forward_rng = nvtx.start_range(message=self.message + "_forward",
                                            color=self.forward_color)
    
    def _forward_post_hook_fn(self, module, input, output):
        nvtx.end_range(self.forward_rng)
    
    def _backward_pre_hook_fn(self, module, grad_input, grad_output):
        self.backward_rng = nvtx.start_range(message=self.message + "_backward",
                                             color=self.backward_color)

    def _backward_post_hook_fn(self, module, grad_input, grad_output):
        nvtx.end_range(self.backward_rng)

        